package com.nanohadoop.yarn;

import com.nanohadoop.mapreduce.JobConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * YARN资源管理器
 */
public class ResourceManager {
    private static final Logger log = LoggerFactory.getLogger(ResourceManager.class);

    private final int port;
    private final String host;
    private final ResourceScheduler scheduler;
    private final ExecutorService threadPool;
    private volatile boolean running;
    private ServerSocket serverSocket;

    public ResourceManager(ResourceScheduler scheduler) {
        this.host = "localhost";
        this.port = 8032;
        this.scheduler = scheduler;
        this.threadPool = Executors.newFixedThreadPool(10);
        this.running = false;
    }

    /**
     * 启动ResourceManager
     */
    public void start() throws IOException {
        if (running) {
            return;
        }
        running = true;

        serverSocket = new ServerSocket(port);
        log.debug("ResourceManager started on " + host + ":" + port);

        // 启动接受客户端连接的线程
        threadPool.submit(this::acceptClientConnections);
    }

    /**
     * 停止ResourceManager
     */
    public void stop() {
        if (!running) {
            return;
        }
        running = false;

        try {
            if (serverSocket != null) {
                serverSocket.close();
            }
            threadPool.shutdown();
        } catch (IOException e) {
            e.printStackTrace();
        }
        log.debug("ResourceManager stopped");
    }

    /**
     * 接受客户端连接
     */
    private void acceptClientConnections() {
        while (running) {
            try {
                Socket clientSocket = serverSocket.accept();
                threadPool.submit(() -> handleClientRequest(clientSocket));
            } catch (IOException e) {
                if (running) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 处理客户端请求
     */
    private void handleClientRequest(Socket clientSocket) {
        // 实际实现中，这里应该处理客户端请求
        // 为简化实现，我们仅打印连接信息
        log.debug("Client connected: " + clientSocket.getInetAddress());
        try {
            clientSocket.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 提交应用程序
     */
    public ApplicationMaster submitApplication(JobConfiguration jobConf) {
        String appId = scheduler.submitApplication(
                jobConf.getJobName(),
                "user",
                new Resource(1024, 1), // 默认1GB内存，1个CPU核心
                jobConf.getNumReducers() + 1 // mapper数量 + reducer数量
        );
        return scheduler.getApplication(appId);
    }

    /**
     * 获取ApplicationMaster
     */
    public ApplicationMaster getApplication(String appId) {
        return scheduler.getApplication(appId);
    }

    /**
     * 获取ResourceManager地址
     */
    public String getAddress() {
        return host + ":" + port;
    }
}